import torch
import numpy as np

def get_TNN1(three_way_tensor):
    tensor_A = torch.fft.fftn(three_way_tensor,s=3,dim=2)
    A_ = tensor_A[:,:,0]
    for i in range(1,3):
        A_ = torch.block_diag(A_, tensor_A[:,:,i])
    # A_ = torch.block_diag(tensor_A[:,:,0], tensor_A[:,:,1], tensor_A[:,:,2])
    # A_ = torch.block_diag(tensor_A[:, :, 0], tensor_A[:, :, 1], tensor_A[:, :, 2], tensor_A[:, :, 3], tensor_A[:, :, 4], tensor_A[:, :, 5])
    _, sigma, _ = torch.svd(A_, some=True)
    nuclear = torch.sum(sigma)
    return nuclear

def get_TNN2(three_way_tensor):
    three_way_tensor = three_way_tensor.numpy()
    bcirc_three_way_tensor = np.block([
        [three_way_tensor[:,:,0], three_way_tensor[:,:,2], three_way_tensor[:,:,1]],
        [three_way_tensor[:,:,1], three_way_tensor[:,:,0], three_way_tensor[:,:,2]],
        [three_way_tensor[:,:,2], three_way_tensor[:,:,1], three_way_tensor[:,:,0]]
    ])
    bcirc_three_way_tensor = torch.from_numpy(bcirc_three_way_tensor)
    _, sigma, _ = torch.svd(bcirc_three_way_tensor, some=True)
    nuclear = torch.sum(sigma)
    return nuclear


three_way_tensor = torch.rand(2,2,3)
print(get_TNN1(three_way_tensor))
print(get_TNN2(three_way_tensor))